from math import sqrt
import time
import cPickle

import pygame
from pygame.locals import *


#######
# All coordinates are in RoboRealm referential, (0,0) is the bottom, left
# pygame uses (0,0) as the top, left
class QuadTree:

    qtree = list()
    n = 0
    coords = list()
    screen = []
    G = dict()

    def __init__( self, lb = [0,0], rt = [1,1]):
        self.qtree = [ 0, lb, rt, [], [], [], [], [], [], [1]]
        self.qtree[8] = sumd2( self.qtree[1], self.qtree[2])
        self.n=0;
        self.coords = [self.qtree[8]]

    def numcells(self):
        return self.n

    def getcellsboundaries(self, cellbond, qt = []):

        if qt == []:
            qt = self.qtree

        if qt[3]== []:
           return [qt[8], qt[1], qt[2]]
        else:
            cellbond.append(self.getcellsboundaries( cellbond, qt[4] ))
            cellbond.append(self.getcellsboundaries( cellbond, qt[7] ))
            cellbond.append(self.getcellsboundaries( cellbond, qt[5] ))
            cellbond.append(self.getcellsboundaries( cellbond, qt[6] ))

        return cellbond


    def updatedist(self):
        # create distances
        self.G = dict()
        for orig in range(0,self.n+1):
            H = dict()
            for dest in range(0,self.n+1):
                if self.freespace( self.getcoordst( orig ), self.getcoordst( dest ) )==False:
                    H[dest] = 1000
                else:
                    H[dest] = dist( self.getcoordst( orig ), self.getcoordst( dest ) )
            self.G[orig] = H



    def setscreen( self, scr):
        self.screen = scr

    def getprop( self, sn, qt = -1):

        if qt == -1:
            qt = self.qtree
            
        if qt[0] == sn:
            return qt[9][0]
        elif qt[0] == -1:
            ret = getprop( self, sn, qt[4] )
            if  ret != -1:
                return ret
            ret = getprop( self, sn, qt[5] )
            if  ret != -1:
                return ret
            ret = getprop( self, sn, qt[6] )
            if  ret != -1:
                return ret
            ret = getprop( self, sn, qt[7] )
            if  ret != -1:
                return ret
        return -1

    def setprop( self, sn, prop, qt = -1 ):

        if qt == -1:
            qt = self.qtree

        if qt[0] == sn:
            qt[9] = [prop]
            return True
        elif qt[0] == -1:
            if self.setprop( sn, prop, qt[4] ):
                return True
            elif self.setprop( sn, prop, qt[5] ):
                return True
            elif self.setprop( sn, prop, qt[6] ):
                return True
            elif self.setprop( sn, prop, qt[7] ):
                return True
            else:
                return False
        else:
            return False




    def drawmqtree( self, qt = [] ):
        if self.screen == []:
            print "no screen, nowhere to draw"
            return

        if qt == []:
            qt = self.qtree

        if qt[3]== []:
            self.plotsq( qt[8], qt[1], qt[2], qt[9][0], qt[0])
        else:
            self.drawmqtree( qt[4] );
            self.drawmqtree( qt[7] );
            self.drawmqtree( qt[5] );
            self.drawmqtree( qt[6] );


    def plotsq( self, cc, bb, tt, prop, num ):
        if prop == 1:
            color = pygame.Color( 0, 255, 0, 1)
        else:
            color = pygame.Color( 255, 0, 0, 1)
    #
        pygame.draw.rect( self.screen, color, pygame.Rect( bb[0], self.qtree[2][1]-tt[1], abs(tt[0]-bb[0]), abs(tt[1]-bb[1])), 5)

        font = pygame.font.Font(None, 20)
        text = font.render( str(num), True, (255, 255, 255), (159, 182, 205))
        # Create a rectangle
        textRect = text.get_rect()
        textRect.centerx = cc[0]
        textRect.centery = self.qtree[2][1]-cc[1]
        # Blit the text
        self.screen.blit(text, textRect)

        return bb


    def looknode( self, T, qt = -1):

        if qt == -1:
            qt = self.qtree

        if qt[3]== []:
            # no cell subdivision, return label
            return qt[0], qt[9]
        else:
            if T[0]<qt[3][0]:
                if T[1]<qt[3][1]:
                    nd, prop = self.looknode( T, qt[4]);
                else:
                    nd, prop = self.looknode( T, qt[7]);
            else:
                if (T[1]<qt[3][1]):
                    nd, prop = self.looknode( T, qt[5]);
                else:
                    nd, prop = self.looknode( T, qt[6]);
        return nd, prop

    
    def freespace( self, a, b):

        xi = float(b[0]-a[0])/50.0
        yi = float(b[1]-a[1])/50.0
        for step in range(0,51):
            x = a[0] + xi * step
            y = a[1] + yi * step

            nd, prop = self.looknode( [x,y] )
            if prop[0]==0:
                return False

        return True


    def getcoordst( self, st):
        return self.coords[st]

    def addnode( self, T, N = [], qt = [], CoordList = []):

        if qt == []:
            N = self.n
            qt = self.qtree
            CoordList = self.coords
        
        if qt[3]== []: #empty slot, add new
            nd = qt[0]
            if( float(prod(sub(qt[2],qt[1])))<=100):
                return [qt, CoordList, N]

            T = sumd2(qt[1],qt[2]);
            qt[3] = T;

            qt[4] = [ qt[0], qt[1], [T[0], T[1]], [], [], [], [], [], [], [1]];
            qt[4][8] = sumd2(qt[4][1],qt[4][2]);
            CoordList[qt[0]] = qt[4][8];
            qt[5] = [ N+1, [T[0], qt[1][1]], [qt[2][0], T[1]], [], [], [], [], [], [], [1] ]
            qt[5][8] = sumd2(qt[5][1],qt[5][2])
            CoordList.append( qt[5][8] );
            qt[6] = [ N+2, [T[0], T[1]], qt[2], [], [], [], [], [], [], [1] ];
            qt[6][8] = sumd2(qt[6][1],qt[6][2])
            CoordList.append( qt[6][8] )
            qt[7] = [ N+3, [qt[1][0], T[1]], [T[0], qt[2][1]], [], [], [], [], [], [], [1] ];
            qt[7][8] = sumd2(qt[7][1],qt[7][2]);
            CoordList.append( qt[7][8] )
            qt[0] = -1;
            qt[8] = [];
            N = N + 3
            self.n = N

        else:
            if T[0]<qt[3][0]:
                if T[1]<qt[3][1]:
                    [nd, CoordList, N] = self.addnode( T, N, qt[4], CoordList);
                    qt[4] = nd;
                else:
                    [ nd, CoordList, N] = self.addnode( T, N, qt[7], CoordList);
                    qt[7] = nd;
            else:
                if T[1]<qt[3][1]:
                    [ nd, CoordList, N] = self.addnode( T, N, qt[5], CoordList);
                    qt[5] = nd;
                else:
                    [ nd, CoordList, N] = self.addnode( T, N, qt[6], CoordList);
                    qt[6] = nd;

        return [qt, CoordList, N]


    def tofile(self, name = 'save.txt'):
        cPickle.dump( [self.qtree, self.n, self.coords, self.G],\
            open( name, 'wb'))

    def fromfile(self, name = 'save.txt'):
        nobj = cPickle.load(open( name, 'rb'))
        self.qtree = nobj[0]
        self.n = nobj[1]
        self.coords = nobj[2]
        self.G = nobj[3]




def prod(a):
    b = a[0]*a[1]
    return b

def sub(a,b):
    c = [ [], []]
    c[0] = a[0]-b[0]
    c[1] = a[1]-b[1]
    return c

def sumd2(a,b):
    c = [ [], []]
    c[0] = (a[0]+b[0])/2
    c[1] = (a[1]+b[1])/2
    return c

def dist(a,b):
    return sqrt((a[0]-b[0])*(a[0]-b[0])+(a[1]-b[1])*(a[1]-b[1]))
